import re
import sys
sys.path.append("/Users/ender/Code_repo/rosalind-brush-questions")
from utils import read_seqfile, read_fasta
import numpy as np


def parser_codon_table(file, stop_symbol="*"):
    codon_table = {}
    with open(file, "r") as f:
        tmp_table_file = "      ".join([line.strip("\n") for line in f.readlines()])
        for codon_aa in re.findall("\S{3}\s\S{1,4}", tmp_table_file):
            c, aa = codon_aa.split(" ")
            if aa == "Stop":
                codon_table[c] = stop_symbol
            else:
                codon_table[c] = aa

    return codon_table


def dna2rna(dna: str) -> str:
    return dna.replace('T', 'U')


def rna2protein(rna: str, codon_table: dict) -> str:
    protein = ""
    for i in range(0, len(rna), 3):
        protein += codon_table[rna[i:i+3]]
    return protein


codon_table = parser_codon_table("./codon_table.txt")

fasta = read_fasta("./rosalind_splc.txt")
dna = fasta.pop(list(fasta.keys())[0])
introns = fasta


introns_idx = np.zeros((len(introns), ), dtype=int)
introns_len = np.zeros((len(introns), ), dtype=int)
for i, (name, intron) in enumerate(introns.items()):
    if intron in dna:
        introns_idx[i] = dna.find(intron)
        introns_len[i] = len(intron)

order = np.argsort(introns_idx)
introns_idx = introns_idx[order]
introns_len = introns_len[order]

s = 0
exons = []
for intron_idx, intron_len in zip(introns_idx, introns_len):
    e = intron_idx
    # print(s, e)
    exons.append(dna[s:e])
    s = e + intron_len
exons.append(dna[s:])


cdna = "".join(exons)
print(len(cdna))
rna = dna2rna(cdna)
protein = rna2protein(rna, codon_table)
print(len(protein))
print(protein)
